#include "PCode.hpp"
#include <iomanip>

//test

const string get_op_name(ins_type op) {
    static const int n_ops = 8; 
    static ins_type ops[n_ops] = {
        INS_LIT, INS_LOD, INS_STO, INS_CAL, 
        INS_INT, INS_JMP, INS_JPC, INS_OPR
    };
    static string names[n_ops] = {
        "LIT", "LOD", "STO", "CAL",
        "INT", "JMP", "JPC", "OPR"
    };
    static unordered_map<ins_type, string> op_name_table;
    static bool initialized = false;
    if(!initialized) {
        for(int i = 0; i < n_ops; i ++) {
            op_name_table[ops[i]] = names[i];
        }
        initialized = true;
    }
    return  op_name_table.find(op) == op_name_table.end() 
            ? "Err" : op_name_table[op];
}

pcode_transformer::pcode_transformer(void) {}
pcode_transformer::pcode_transformer(const SymbolTable& symbol_table) {
    this->symbol_table = symbol_table;
}

int pcode_transformer::fill_addr(int idx_from, int idx_to, int addr, ins_type type) {

    int total = 0;
    if(type != INS_JPC && type != INS_JMP)
        return 0;

    for(int i = idx_from; i <= idx_to; i ++) {
        if(this->target_codelines[i].OP == type 
        && this->target_codelines[i].need_fill) {
            this->target_codelines[i].d = addr;
            this->target_codelines[i].need_fill = false;
            total++;
        }
    }
    return total;
}

void pcode_transformer::set_symbol_table(const SymbolTable& symbol_table) {
    this->symbol_table = symbol_table;
}

void pcode_transformer::clear_instruction(void) {
    this->target_codelines.clear();
}

void pcode_transformer::insert_instruction(const st_ins& ins) {
    this->target_codelines.push_back(ins);
}

void pcode_transformer::insert_instruction(ins_type op, int l, int d) {
    this->target_codelines.push_back(st_ins(op, l, d));
}

pcode_transformer& pcode_transformer::operator<<(const st_ins& ins) {
    this->target_codelines.push_back(ins);
    return (*this);
}

std::set<string> pcode_transformer::read_AST(const TreeDefine* root, int cur_level, int last_node_type) {

    int node_type;
    int expr_type;
    std::set<string> vars;

    if(!root) {
        return vars;
    }
    node_type = root->getNodeType();
    expr_type = root->getExpType();
    switch(node_type) 
    {
    case PROGRAM: case PROGRAMHEAD: 
        return gen_head(root, cur_level+1, node_type);
    case FUNCTIONHEAD: case PROCEDUREHEAD:
        return gen_head(root, cur_level+1, node_type);
    case PROGRAMBODY:
        return gen_programbody(root, cur_level+1, node_type);
    case DK_VAR:
        return gen_dk_var(root, cur_level, node_type);
    case EK_CONST:
        return gen_ek_const(root, cur_level, node_type);;
    case OP:
        return gen_op(root, cur_level, node_type);;
    case FUNC_ID:
        return gen_func_id(root, cur_level, node_type);;
    case FUNC_SYS:
        return gen_func_sys(root, cur_level, node_type);;
    case EK_ID:
        return gen_ek_id(root, cur_level, node_type);
    case TO:
        return gen_to(root, cur_level, node_type);;
    case DOWNTO:
        return gen_downto(root, cur_level, node_type);;
    case READ:
        return gen_read(root, cur_level, node_type);;
    case WRITE:
        return gen_write(root, cur_level, node_type);;
    case LT:
        return gen_lt(root, cur_level, node_type);;
    case LE:
        return gen_le(root, cur_level, node_type);;
    case GT:
        return gen_gt(root, cur_level, node_type);;
    case GE:
        return gen_ge(root, cur_level, node_type);;
    case EQUAL:
        return gen_equal(root, cur_level, node_type);;
    case UNEQUAL:
        return gen_unequal(root, cur_level, node_type);;
    case NOT:
        return gen_not(root, cur_level, node_type);;
    case OR:
        return gen_or(root, cur_level, node_type);;
    case AND:
        return gen_and(root, cur_level, node_type);;
    case ASSIGN:
        return gen_assign(root, cur_level, node_type);;
    case IF:
        return gen_if(root, cur_level, node_type);;
    case REPEAT:
        return gen_repeat(root, cur_level, node_type);;
    case WHILE:
        return gen_while(root, cur_level, node_type);;
    case FOR:
        return gen_for(root, cur_level, node_type);;
    case EK_CASE:
        return gen_ek_case(root, cur_level, node_type);;
    case SK_CASE:
        return gen_sk_case(root, cur_level, node_type);;
    case PROC_ID:
        return gen_proc_id(root, cur_level, node_type);;
    case PROC_SYS:
        return gen_proc_sys(root, cur_level, node_type);;
    case MUL:
        return gen_mul(root, cur_level, node_type);
    case DIV:
        return gen_div(root, cur_level, node_type);
    case PLUS:
        return gen_plus(root, cur_level, node_type);
        //return gen_positive(root, cur_level, node_type);
    case MINUS:
        if(root->getChildren().size() == 1){
            return gen_negative(root, cur_level, node_type);
        }
        else{
            return gen_minus(root, cur_level, node_type);
        }   
    case MOD:
        return gen_mod(root, cur_level, node_type);
    default:
        break;
    }
    return vars;
}

std::set<string> pcode_transformer::gen_head(const TreeDefine* root, int cur_level, int last_node_type) {
    
    // 在程序头/函数头等中，并不包含实际的操作，也无法转换成指令
    // auto children = root->getChildren();
    // for(auto iter = children.begin();
    //     iter != children.end(); iter++) {
    //     read_AST(*iter, cur_level, last_node_type);
    // }
    return std::set<string>();
}

std::set<string> pcode_transformer::gen_programbody(const TreeDefine* root, int cur_level, int last_node_type) {

    // 分配的栈空间
    std::set<string> vars;
    auto children = root->getChildren();
    for(auto iter = children.begin();
        iter != children.end(); iter++) {
        auto part_vars = read_AST(*iter, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    if(vars.size() != 0) {
        (*this) << st_ins(INS_INT, 0, vars.size());
    }
    return vars;
}

std::set<string> pcode_transformer::gen_ek_id(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    auto ret = SymbolTable::lookupVar(root->getAttribute()._string,cur_level);
    
    // todo id array record
    // (*this) << st_ins(INS_LOD, , ret.memPos.offset);
    return vars;
}

std::set<string> pcode_transformer::gen_dk_var(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    
    return vars;
}

std::set<string> pcode_transformer::gen_mul(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }

    (*this) << st_ins(INS_OPR, 0, AL_MUL);
    return vars;
}

std::set<string> pcode_transformer::gen_div(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }

    (*this) << st_ins(INS_OPR, 0, AL_DIV);
    return vars;
}

std::set<string> pcode_transformer::gen_plus(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }

    (*this) << st_ins(INS_OPR, 0, AL_ADD);
    return vars;
}

//std::set<string> pcode_transformer::gen_positive(const TreeDefine* root, int cur_level, int last_node_type) {
//
//
//}

std::set<string> pcode_transformer::gen_minus(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;  
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }

    (*this) << st_ins(INS_OPR, 0, AL_SUB);
    return vars;
}

std::set<string> pcode_transformer::gen_negative(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }

    (*this) << st_ins(INS_OPR, 0, AL_NEG);
    return vars;
}

std::set<string> pcode_transformer::gen_mod(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }

    (*this) << st_ins(INS_OPR, 0, AL_MOD);
    return vars;

}

//std::set<string> pcode_transformer::gen_op(const TreeDefine* root, int cur_level, int last_node_type) {
//    auto items = root->getChildren();
//    std::set<string> vars;
//
//    return vars;
//}
std::set<string> pcode_transformer::gen_ek_const(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    return vars;
}
std::set<string> pcode_transformer::gen_func_id(const TreeDefine* root, int cur_level, int last_node_type) {
    // TODO 不知道长什么样子
    auto items = root->getChildren();
    std::set<string> vars;
    
    return vars;
}
std::set<string> pcode_transformer::gen_func_sys(const TreeDefine* root, int cur_level, int last_node_type) {
    // TODO 不知道长什么样子
    auto items = root->getChildren();
    std::set<string> vars;

    return vars;
}
//std::set<string> pcode_transformer::gen_to(const TreeDefine* root, int cur_level, int last_node_type) {
//    auto items = root->getChildren();
//    std::set<string> vars;
//
//    return vars;
//}
//std::set<string> pcode_transformer::gen_downto(const TreeDefine* root, int cur_level, int last_node_type) {
//    auto items = root->getChildren();
//    std::set<string> vars;
//
//    return vars;
//}
std::set<string> pcode_transformer::gen_read(const TreeDefine* root, int cur_level, int last_node_type) {

    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
        (*this) << st_ins(INS_OPR, 0, AL_SCN);
    }
    return vars;
}
std::set<string> pcode_transformer::gen_write(const TreeDefine* root, int cur_level, int last_node_type) {

    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
        (*this) << st_ins(INS_OPR, 0, AL_PRT);
    }
    return vars;
}
std::set<string> pcode_transformer::gen_lt(const TreeDefine* root, int cur_level, int last_node_type) {

    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    (*this) << st_ins(INS_OPR, 0, AL_LES);
    return vars;
}
std::set<string> pcode_transformer::gen_le(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    (*this) << st_ins(INS_OPR, 0, AL_LEQ);
    return vars;
}
std::set<string> pcode_transformer::gen_gt(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    (*this) << st_ins(INS_OPR, 0, AL_GTR);
    return vars;
}
std::set<string> pcode_transformer::gen_ge(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    (*this) << st_ins(INS_OPR, 0, AL_GEQ);
    return vars;
}
std::set<string> pcode_transformer::gen_equal(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    (*this) << st_ins(INS_OPR, 0, AL_EQU);
    return vars;
}
std::set<string> pcode_transformer::gen_unequal(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    (*this) << st_ins(INS_OPR, 0, AL_NEQ);
    return vars;
}
std::set<string> pcode_transformer::gen_not(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    (*this) << st_ins(INS_OPR, 0, AL_NOT);
    return vars;
}
std::set<string> pcode_transformer::gen_or(const TreeDefine* root, int cur_level, int last_node_type) {

    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        (*this) << st_ins(INS_JPC, 0, 0, true);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    return vars;
}
std::set<string> pcode_transformer::gen_and(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        (*this) << st_ins(INS_OPR, 0, AL_NOT);
        (*this) << st_ins(INS_JPC, 0, 0, true);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    return vars;
}
std::set<string> pcode_transformer::gen_assign(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    string left_val_name;
    string attri_name;
    int array_idx;

    for(auto item=items.begin();
        item != items.end(); item++) {
        auto part_vars = read_AST(*item, cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    auto var_ID_node = items[0]->getChildren()[0];
    auto var_type_node = var_ID_node->getChildren();
    left_val_name = var_ID_node->getAttribute()._string;
    // TODO: 目前只实现了三种简单的查找变量，待完善
    if(!var_type_node.empty()) {
        auto IdxOrAttri_node = var_type_node[0]->getChildren()[0];
        // 结构体
        if(var_type_node[0]->getNodeType() == TK_RECORD) {
            attri_name = IdxOrAttri_node->getChildren()[0]->getAttribute()._string;
            auto res = SymbolTable::lookupRecord(left_val_name, attri_name, cur_level);
            (*this) << st_ins(INS_STO, res.jumpLevel, res.totalOff);
        } else {    // 数组
            array_idx = IdxOrAttri_node->getChildren()[0]->getAttribute()._int;
            auto res = SymbolTable::lookupVar();
            // TODO: 目前的实现是基于“数组的起始下标为0”的假设，待完善
            (*this) << st_ins(INS_STO, res.memPos.basePos, res.memPos.offset+array_idx); 
        }
    } else {    // 简单变量
        auto res = SymbolTable::lookupVar(left_val_name, cur_level);
        (*this) << st_ins(INS_STO, res.memPos.basePos, res.memPos.offset); 
    }
    
    return vars;
}
std::set<string> pcode_transformer::gen_if(const TreeDefine* root, int cur_level, int last_node_type) {

    auto items = root->getChildren();
    int items_cnt = items.size();
    std::set<string> vars;
    int begin_pos = this->target_codelines.size();
    int if_then_pos=-1;
    int else_pos=-1;
    int end_pos;
    for(int i = 0; i < items_cnt; i ++) {
        auto part_vars = read_AST(items[i], cur_level, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
        if(i == 0) {
            if_then_pos = this->target_codelines.size();
        } else if(i == 1 && i < items_cnt-1) {
            (*this) << st_ins(INS_JMP, 0, 0, true);         // # x
            else_pos = this->target_codelines.size();
        }
    }
    end_pos = this->target_codelines.size();
    
    // 回填到if语句到then之前的指令中
    fill_addr(begin_pos, if_then_pos-1, else_pos, INS_JPC);
    // else语句的前一条即是JMP指令
    if(items_cnt == 3) {
        fill_addr(else_pos-1, else_pos-1, end_pos, INS_JMP);    // 回填上方x处的地址
    }

    return vars;
}
std::set<string> pcode_transformer::gen_repeat(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    
    int repeat_begin_pos = this->target_codelines.size()+1; //指向repeat开始的位置

    auto part_vars = read_AST(items[0], cur_level, last_node_type);
    vars.insert(part_vars.begin(), part_vars.end());

    auto part_vars = read_AST(items[1], cur_level, last_node_type);
    vars.insert(part_vars.begin(), part_vars.end());

    (*this) << st_ins(INS_JPC, 0, repeat_begin_pos);

    return vars;
}
std::set<string> pcode_transformer::gen_while(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    int while_begin_pos = this->target_codelines.size()+1; //指向while开始的位置

    auto part_vars = read_AST(items[0], cur_level, last_node_type);
    vars.insert(part_vars.begin(), part_vars.end());

    (*this) << st_ins(INS_JPC, 0, 0, true);

    auto part_vars = read_AST(items[1], cur_level, last_node_type);
    vars.insert(part_vars.begin(), part_vars.end());

    (*this) << st_ins(INS_JMP, 0, while_begin_pos);

    int while_end_pos = this->target_codelines.size(); //指向while结束的位置

    fill_addr(while_begin_pos, while_end_pos, while_end_pos, INS_JPC);
    return vars;
}
std::set<string> pcode_transformer::gen_for(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    string id_name = items[0]->getAttribute()._string;
    auto i_ret = SymbolTable::lookupVar(id_name, cur_level);

    auto part_vars = read_AST(items[0], cur_level, last_node_type);
    vars.insert(part_vars.begin(), part_vars.end());
    auto part_vars = read_AST(items[1], cur_level, last_node_type);
    vars.insert(part_vars.begin(), part_vars.end());

    auto ret = SymbolTable::lookupVar(root->getAttribute()._string,cur_level);
    (*this) << st_ins(INS_STO, ret.memPos.basePos, ret.memPos.offset);

    auto part_vars = read_AST(items[2], cur_level, last_node_type);
    vars.insert(part_vars.begin(), part_vars.end());

    (*this) << st_ins(INS_OPR, 0, AL_LEQ);
    int comp_pos = this->target_codelines.size();
    (*this) << st_ins(INS_JPC, 0, 0, true);

    auto part_vars = read_AST(items[3], cur_level, last_node_type);
    vars.insert(part_vars.begin(), part_vars.end());


    (*this) << st_ins(INS_LOD, i_ret.memPos.basePos, i_ret.memPos.offset); // 取出 i
    (*this) << st_ins(INS_LIT, 0, 1);
    (*this) << st_ins(INS_OPR, 0, AL_ADD);
    (*this) << st_ins(INS_STO, 0, i_ret.memPos.basePos, i_ret.memPos.offset);
    (*this) << st_ins(INS_JMP, 0, comp_pos);

    int do_end_pos = this->target_codelines.size();
    fill_addr(comp_pos, do_end_pos, do_end_pos, INS_JPC);
    
    return vars;
}
//std::set<string> pcode_transformer::gen_ek_case(const TreeDefine* root, int cur_level, int last_node_type) {
//    auto items = root->getChildren();
//    std::set<string> vars;
//
//    return vars;
//}
//std::set<string> pcode_transformer::gen_sk_case(const TreeDefine* root, int cur_level, int last_node_type) {
//    auto items = root->getChildren();
//    std::set<string> vars;
//
//    return vars;
//}
std::set<string> pcode_transformer::gen_proc_id(const TreeDefine* root, int cur_level, int last_node_type) {
    // TODO 还需要加入INT
    auto items = root->getChildren();
    std::set<string> vars;
    for(auto item=items.begin();item!=items.end();item++){
        auto part_vars = read_AST(*item, cur_level+1, last_node_type);
        vars.insert(part_vars.begin(), part_vars.end());
    }
    return vars;
}
std::set<string> pcode_transformer::gen_proc_sys(const TreeDefine* root, int cur_level, int last_node_type) {
    auto items = root->getChildren();
    std::set<string> vars;
    if(items[0]->getAttribute()._int == WRITE){
        auto children = items[0]->getChildren();
        for(auto child=children.begin();child!=children.end();child++){
            auto part_vars = read_AST(*child, cur_level+1, last_node_type);
            vars.insert(part_vars.begin(), part_vars.end());
            (*this) << st_ins(INS_OPR, 0, AL_PRT);
        }
    }
    return vars;
}

void pcode_transformer::operator<<(const TreeDefine* AST) {
    
    read_AST(AST, 0, 0);
}

ostream& operator<<(ostream& output, const pcode_transformer& pt) {

    const int begin_addr = 1;
    vector<st_ins> ins = pt.target_codelines;
    int n_codeline = ins.size();

    output << "+------+-----+-----+-----+\n";
    output << "| addr | OP  |  L  |  d  |\n";
    output << "+------+-----+-----+-----+\n";
    for(int i = 0; i < n_codeline; i ++) {
        output << "| " << setw(4) << right << i+begin_addr;
        output << " | "<< setw(3) << left << get_op_name(ins[i].OP);
        output << " | "<< setw(3) << left << ins[i].L;
        output << " | "<< setw(3) << left << ins[i].d << " |\n";
        output << "+------+-----+-----+-----+\n";
    }
    return output;
}